package com.smart.speaker.websocket.chat;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationOutput;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.tools.ToolCallBase;
import com.smart.speaker.config.ModelConfig;
import com.smart.speaker.util.PromptUtils;
import com.smart.speaker.util.ToolUtils;
import com.smart.speaker.websocket.ChatHandler;
import com.smart.speaker.websocket.SessionState;
import io.reactivex.Flowable;
import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

@Slf4j
public class Normal extends Chat {

    private final ChatHandler chatHandler;
    private final Generation generation;
    private final List<Message> historyList;
    private final SessionState sessionState;

    public Normal(Generation generation, SessionState sessionState, ChatHandler chatHandler) {
        this.chatHandler = chatHandler;
        this.generation = generation;
        this.sessionState = sessionState;
        this.historyList = new LinkedList<>();
    }

    @Override
    public void chat(String req) {
        try {
            Message systemMsg = Message.builder().role(Role.SYSTEM.getValue())
                    .content(PromptUtils.chat())
                    .build();
            Message userMsg = Message.builder().role(Role.USER.getValue()).content(req).build();
            var param = GenerationParam.builder()
                    .apiKey(ModelConfig.Key.apiKey)
                    .model(ModelConfig.Key.model)
                    .messages(getChatHistory(systemMsg, userMsg))
                    .tools(ToolUtils.getDefaultToolList())
                    .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                    .incrementalOutput(true)
                    .build();
            Flowable<GenerationResult> flowable = generation.streamCall(param);
            List<ToolCallBase> tools = new ArrayList<>();
            flowable.blockingForEach(message -> {
                log.info(message.toString());
                List<GenerationOutput.Choice> choiceList = message.getOutput().getChoices();
                for (GenerationOutput.Choice choice : choiceList) {
                    if (choice.getMessage().getToolCalls() == null) {
                        String content = choice.getMessage().getContent();
                        sessionState.appendResponse(content);
                        String sentence = sessionState.getSentence();
                        chatHandler.synthesizer(sentence);
                    } else {
                        tools.addAll(choice.getMessage().getToolCalls());
                    }
                }
            });
            String sentence = sessionState.getLastSentence();
            chatHandler.synthesizer(sentence);
            chatHandler.executeTool(tools);
            String response = sessionState.getResponse();
            if (response != null) {
                addHistory(userMsg, Message.builder().role(Role.ASSISTANT.getValue()).content(response).build());
            }
        } catch (Exception e) {
            log.error("调用大模型失败", e);
        }
    }

    private List<Message> getChatHistory(Message systemMsg, Message userMsg) {
        while (historyList.size() > 10) {
            historyList.removeFirst();
        }
        if (historyList.size() % 2 != 0) {
            historyList.removeFirst();
        }
        List<Message> messages = new ArrayList<>(historyList.size() + 2);
        messages.add(systemMsg);
        messages.addAll(historyList);
        messages.add(userMsg);
        return messages;
    }

    private void addHistory(Message userMsg, Message aiMsg) {
        historyList.add(userMsg);
        historyList.add(aiMsg);
    }
}
